package forward

import (
	"net"
	"sync"
	"time"

	"gitee.com/sunny-zhangqiang/zcache/comm/slru"
	"golang.org/x/sync/singleflight"
)

const (
	defaultDNSResultTTL = 180 // 3 min
)

var (
	resolver *DNSResolver
)

// FIXME: 这是一个简单的DNS解析器, 后续应该按需定制.

type DNSResult struct {
	Address    []string
	ExpireTime int64
}

type DNSResolver struct {
	cacheMux sync.Mutex
	cache    *slru.S4LRU
	g        singleflight.Group
}

func (r *DNSResolver) lookupInCache(host string) *DNSResult {
	r.cacheMux.Lock()
	defer r.cacheMux.Unlock()

	v, ok := r.cache.Get(host)
	if !ok {
		return nil
	}

	res := v.(*DNSResult)
	now := time.Now().Unix()
	if now > res.ExpireTime {
		r.cache.Remove(host)
		return nil
	}

	return v.(*DNSResult)
}

func (r *DNSResolver) setInCache(host string, res *DNSResult) {
	r.cacheMux.Lock()
	defer r.cacheMux.Unlock()
	r.cache.Set(host, res)
}

func (r *DNSResolver) doLookup(host string) (*DNSResult, error) {
	v, err, _ := r.g.Do(host, func() (interface{}, error) {
		return net.LookupAddr(host)
	})

	if err != nil {
		return nil, err
	}
	return &DNSResult{Address: v.([]string)}, nil
}

func (r *DNSResolver) Lookup(host string) (*DNSResult, error) {
	res := r.lookupInCache(host)
	if res != nil {
		return res, nil
	}

	res, err := r.doLookup(host)
	if err != nil {
		return nil, err
	}

	now := time.Now().Unix()
	res.ExpireTime = now + defaultDNSResultTTL
	r.setInCache(host, res)
	return res, nil
}

func init() {
	resolver = &DNSResolver{
		cache: slru.New(128, 0, nil),
	}
}
